/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 */
#ifndef SEND_RECV_EXECUTOR_H
#define SEND_RECV_EXECUTOR_H
#include "dtype_common.h"
#include "hccl_common.h"
#include "adapter_hccp.h"
#include "stream_pub.h"
#include "dispatcher.h"
#include "typical_qp_manager.h"
#include "typical_window_mem.h"
#include "typical_sync_mem.h"

namespace hccl {
constexpr u64 HCCL_CHUNK_SIZE = 1024 * 1024 * 1024;
class SendRecvExecutor {
public:
    SendRecvExecutor(HcclRtStream stream, QpHandle qpHandle,
    const struct mr_info& localWindowMem, const struct mr_info& remoteWindowMem,
    const struct mr_info& localSyncMemPrepare, const struct mr_info& localSyncMemDone, const struct mr_info& localSyncMemAck,
    const struct mr_info& remoteSyncMemPrepare, const struct mr_info& remoteSyncMemDone, const struct mr_info& remoteSyncMemAck,
    u32 immData, const u64 chunkNum = HCCL_CHUNK_SIZE);
    ~SendRecvExecutor();
    HcclResult Init();
    HcclResult Send(void* inputPtr, u64 count, HcclDataType dataType);
    HcclResult Receive(void* outputPtr, u64 count, HcclDataType dataType);
private:
    HcclResult PollCq();
    HcclResult IsOverlappedWithWinMem(void* inputPtr, u64 userMemSize, bool& isOverlapped);
    HcclResult SendRun(DeviceMem& sendBuffer);
    HcclResult ReceiveRun(DeviceMem& receiveBuffer);
    HcclResult ReceiveRunByPollCq(DeviceMem& receiveBuffer);
    HcclResult RecordNotify(void *dstMemPtr, u32 rkey, const void *srcMemPtr, u32 lkey, u64 srcMemSize);
    HcclResult WaitSignal(HcclRtSignal signal);
    HcclResult PayLoad(const void *src, u64 dstOffset, u64 len);
    HcclResult RdmaSendAsync(struct send_wr_v2 &wr);
    HcclResult MemcpyAsyncD2D(hccl::DeviceMem &dst, const hccl::DeviceMem &src, hccl::Stream &stream);
private:
    HcclRtStream stream_;
    QpHandle qpHandle_;
    struct mr_info localWindowMem_;
    struct mr_info remoteWindowMem_;

    struct mr_info localSyncMemPrepare_;
    struct mr_info localSyncMemDone_;
    struct mr_info localSyncMemAck_;

    struct mr_info remoteSyncMemPrepare_;
    struct mr_info remoteSyncMemDone_;
    struct mr_info remoteSyncMemAck_;

    struct mr_info notifySrcMem_;

    HcclRtSignal prepareNotify_ = nullptr;
    HcclRtSignal ackNotify_ = nullptr;
    HcclRtSignal doneNotify_ = nullptr;
    u32 notifySize_;

    u32 immData_ = 0;
    u64 chunkSize_;
    SyncMode notifyWaitMode_;
};
}  // namespace hccl

#endif
